import os
from typing import Generator
from openai import OpenAI

PROMPT = """# 要求：
你将看到一段待检索的文本，请根据文本内容生成一个查询。这些查询应该围绕文本的主要内容、关键词或主题展开，并且要体现一定的多样性和随机性。
生成的查询可以针对文本的不同侧面，例如核心概念、相关背景、或潜在问题。请生成查询时尽量简洁、相关且具备查询意图。此外，生成的查询控制在5到30字之间。
          
# 待检索的文本
{content}

# 请生成查询：
“Query=<你输出的查询>”，其中的“Query=”也是必须回复的；注意，只能有一个查询，只能有一个查询，只能有一个查询。
"""


def traverse_dir(path, tails, file_info_ls=None):
    """

    :param path:
    :param tails:
    :param file_info_ls:
    :return: [(文件路径1，文件名1),(文件路径2，文件名2),(文件路径3，文件名3),...,(文件路径n，文件名n)]
    注意：如路经太深python就无法检测该文件了，遇到这种情况请开启 clean_path功能；
    参考： https://blog.csdn.net/baidu_34928905/article/details/131935428
    """
    path = path
    if file_info_ls is None:
        file_info_ls = set()
    for file in os.listdir(path):
        file_path = os.path.join(path, file)
        if os.path.isdir(file_path):
            traverse_dir(file_path, tails, file_info_ls)
        else:
            file_e_ls = os.path.basename(file_path).split('.')
            file_tail = file_e_ls[-1]
            if (tails is None or file_tail in tails) and not os.path.basename(file_path).startswith('~$'):
                item = (file_path, '.'.join(file_e_ls[0:-1]))
                file_info_ls.add(item)
    return file_info_ls


class Model:

    def __init__(self):
        self.client = OpenAI(api_key="EMPTY", base_url="http://172.16.3.116:2333/v1")
        self.model = "Qwen2-72B-Instruct-AWQ"

    def generate(self, messages, temperature=0.3) -> str:
        res = self.client.chat.completions.create(
            model=self.model,
            messages=messages,
            temperature=temperature,
            stream=False
        )

        return res.choices[0].message.content

    def stream(self, messages, temperature=0.3) -> Generator:

        res = self.client.chat.completions.create(
            model=self.model,
            messages=messages,
            temperature=temperature,
            stream=True
        )

        for chunk in res:
            content = chunk.choices[0].delta.content
            if content is not None:
                yield content
